{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42145ff5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['OMP_NUM_THREADS'] = '4'\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import openai\n",
    "openai.api_key = os.environ['OPENAI_API_KEY']\n",
    "openai.api_base = os.environ.get(\"OPENAI_API_BASE\", \"https://api.openai.com/v1\")\n",
    "from data.serialize import SerializerSettings\n",
    "from models.utils import grid_iter\n",
    "from models.promptcast import get_promptcast_predictions_data\n",
    "from models.darts import get_arima_predictions_data\n",
    "from models.llmtime import get_llmtime_predictions_data\n",
    "from data.small_context import get_datasets\n",
    "from models.validation_likelihood_tuning import get_autotuned_predictions_data\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "def plot_preds(train, test, pred_dict, model_name, show_samples=False):\n",
    "    pred = pred_dict['median']\n",
    "    pred = pd.Series(pred, index=test.index)\n",
    "    plt.figure(figsize=(8, 6), dpi=100)\n",
    "    plt.plot(train)\n",
    "    plt.plot(test, label='Truth', color='black')\n",
    "    plt.plot(pred, label=model_name, color='purple')\n",
    "    # shade 90% confidence interval\n",
    "    samples = pred_dict['samples']\n",
    "    lower = np.quantile(samples, 0.05, axis=0)\n",
    "    upper = np.quantile(samples, 0.95, axis=0)\n",
    "    plt.fill_between(pred.index, lower, upper, alpha=0.3, color='purple')\n",
    "    if show_samples:\n",
    "        samples = pred_dict['samples']\n",
    "        # convert df to numpy array\n",
    "        samples = samples.values if isinstance(samples, pd.DataFrame) else samples\n",
    "        for i in range(min(10, samples.shape[0])):\n",
    "            plt.plot(pred.index, samples[i], color='purple', alpha=0.3, linewidth=1)\n",
    "    plt.legend(loc='upper left')\n",
    "    if 'NLL/D' in pred_dict:\n",
    "        nll = pred_dict['NLL/D']\n",
    "        if nll is not None:\n",
    "            plt.text(0.03, 0.85, f'NLL/D: {nll:.2f}', transform=plt.gca().transAxes, bbox=dict(facecolor='white', alpha=0.5))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1341ded",
   "metadata": {},
   "source": [
    "## Define models ##"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bdf1f3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt4_hypers = dict(\n",
    "    alpha=0.3,\n",
    "    basic=True,\n",
    "    temp=1.0,\n",
    "    top_p=0.8,\n",
    "    settings=SerializerSettings(base=10, prec=3, signed=True, time_sep=', ', bit_sep='', minus_sign='-')\n",
    ")\n",
    "\n",
    "gpt3_hypers = dict(\n",
    "    temp=0.7,\n",
    "    alpha=0.95,\n",
    "    beta=0.3,\n",
    "    basic=False,\n",
    "    settings=SerializerSettings(base=10, prec=3, signed=True, half_bin_correction=True)\n",
    ")\n",
    "\n",
    "\n",
    "promptcast_hypers = dict(\n",
    "    temp=0.7,\n",
    "    settings=SerializerSettings(base=10, prec=0, signed=True, \n",
    "                                time_sep=', ',\n",
    "                                bit_sep='',\n",
    "                                plus_sign='',\n",
    "                                minus_sign='-',\n",
    "                                half_bin_correction=False,\n",
    "                                decimal_point='')\n",
    ")\n",
    "\n",
    "arima_hypers = dict(p=[12,30], d=[1,2], q=[0])\n",
    "\n",
    "model_hypers = {\n",
    "    'LLMTime GPT-3.5': {'model': 'gpt-3.5-turbo-instruct', **gpt3_hypers},\n",
    "    'LLMTime GPT-4': {'model': 'gpt-4', **gpt4_hypers},\n",
    "    'LLMTime GPT-3': {'model': 'text-davinci-003', **gpt3_hypers},\n",
    "    'PromptCast GPT-3': {'model': 'text-davinci-003', **promptcast_hypers},\n",
    "    'ARIMA': arima_hypers,\n",
    "    \n",
    "}\n",
    "\n",
    "model_predict_fns = {\n",
    "    'LLMTime GPT-3': get_llmtime_predictions_data,\n",
    "    'LLMTime GPT-4': get_llmtime_predictions_data,\n",
    "    'PromptCast GPT-3': get_promptcast_predictions_data,\n",
    "    'ARIMA': get_arima_predictions_data,\n",
    "}\n",
    "\n",
    "model_names = list(model_predict_fns.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2da9a874",
   "metadata": {},
   "source": [
    "## Running LLMTime and Visualizing Results ##"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebc3ecb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = get_datasets()\n",
    "ds_name = 'AirPassengersDataset'\n",
    "\n",
    "data = datasets[ds_name]\n",
    "train, test = data # or change to your own data\n",
    "out = {}\n",
    "for model in model_names: # GPT-4 takes a about a minute to run\n",
    "    model_hypers[model].update({'dataset_name': ds_name}) # for promptcast\n",
    "    hypers = list(grid_iter(model_hypers[model]))\n",
    "    num_samples = 10\n",
    "    pred_dict = get_autotuned_predictions_data(train, test, hypers, num_samples, model_predict_fns[model], verbose=False, parallel=False)\n",
    "    out[model] = pred_dict\n",
    "    plot_preds(train, test, pred_dict, model, show_samples=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fspace",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  },
  "vscode": {
   "interpreter": {
    "hash": "9436057e92285046d415c34e216bd357b01decd87fa7e06f42744a4b160880c3"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
